import numpy as np
from math import sqrt
import operator as opt


def normData(dataSet):
    maxVals = dataSet.max(axis=0)
    minVals = dataSet.min(axis=0)
    ranges = maxVals - minVals
    retData = (dataSet - minVals) / ranges
    return retData, ranges, minVals


def kNN(dataSet, labels, testData, k):
    distSquareMat = (dataSet - testData) ** 2
    distSquareSums = distSquareMat.sum(axis=1)
    distances = distSquareSums ** 0.5
    sortedIndices = distances.argsort()
    indices = sortedIndices[:k]
    labelCount = {}
    for i in indices:
        label = labels[i]
        labelCount[label] = labelCount.get(label, 0) + 1
    sortedCount = sorted(labelCount.items(),
                         key=opt.itemgetter(1), reverse=True)
    return sortedCount[0][0]


if __name__ == "__main__":
    dataSet = np.array([[2, 3], [6, 8]])
normDataSet, ranges, minVals = normData(dataSet)
labels = ['a', 'b']
testData = np.array([3.9, 5.5])
normTestData = (testData - minVals) / ranges
result = kNN(normDataSet, labels, normTestData, 1)
print(result)
